import logging
import os

import pandas as pd
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec

from absl import app

import statsmodels.api as sm


def setup_plt():
  # plt.style.use('fivethirtyeight')
  # plt.style.use('seaborn-talk')
  # plt.rcParams['figure.dpi'] = 72
  plt.rcParams['savefig.dpi'] = 100
  plt.rcParams['font.family'] = ['Inconsolata']
  plt.rcParams['font.sans-serif'] = ['Inconsolata']
  plt.rcParams['font.monospace'] = ['Inconsolata']
  # plt.rcParams['figure.constrained_layout.use'] = True

  plt.rcParams['lines.linewidth'] = 0.5
  plt.rcParams['xtick.labelsize'] = 7
  plt.rcParams['ytick.labelsize'] = 7
  '''
  plt.rcParams['lines.linewidth'] = 0.2
  plt.rcParams['grid.linewidth'] = 0.1
  plt.rcParams['font.size'] = 4
  plt.rcParams['legend.fontsize'] = 4
  plt.rcParams['savefig.facecolor'] = '#ffffff'
  '''


def plot_bucket(Y_pred, Y_actual, fig, gs=None, nbuckets=100):
  df = pd.DataFrame({'Y_pred': Y_pred, 'Y_actual': Y_actual})
  df.dropna(inplace=True)
  df.sort_values('Y_pred', inplace=True)

  Y_pred_mid = []
  Y_mean = []
  Y_std = []

  size = df.shape[0]
  for bucket_idx in range(nbuckets):
    from_idx = int(bucket_idx / nbuckets * size)
    to_idx = int((bucket_idx + 1) / nbuckets * size)
    Y_bucket = df.iloc[from_idx:to_idx, 1]
    Y_pred_mid.append(df.iloc[(from_idx + to_idx) // 2, 0])
    Y_mean.append(Y_bucket.mean())
    Y_std.append(Y_bucket.std())

  if gs is None:
    gs = gridspec.GridSpec(6, 1, figure=fig)
  else:
    gs = gs.subgridspec(6, 1)

  ax11 = fig.add_subplot(gs[0:2, :])
  ax11.set_xlabel('Y-pred')
  ax11.set_ylabel('Y-mean')
  ax11.plot(Y_pred_mid, Y_mean)

  ax12 = fig.add_subplot(gs[2, :])
  ax12.set_xlabel('Y-pred')
  ax12.set_ylabel('Y-std')
  ax12.plot(Y_pred_mid, Y_std)

  ax21 = fig.add_subplot(gs[3:5, :])
  ax21.set_xlabel('Y-pred %')
  ax21.set_ylabel('Y-mean')
  ax21.plot(Y_mean)

  ax22 = fig.add_subplot(gs[5, :])
  ax22.set_xlabel('Y-pred %')
  ax22.set_ylabel('Y-std')
  ax22.plot(Y_std)


def plot_timed_y(timed_Y, time_tick_labels, fig, gs=None, log=False, title=None):
  if gs is None:
    gs = gridspec.GridSpec(1, 1, figure=fig)
  else:
    gs = gs.subgridspec(1, 1)

  ax = fig.add_subplot(gs[:, :])
  if log:
    ax.plot(np.log10(np.asarray(timed_Y)), '.-')
  else:
    ax.plot(timed_Y, '.-')
  plt.axvline(x=len(time_tick_labels) // 2, color='black', alpha=0.3, linewidth=0.5)

  ax.set_ylabel('Y-mean')
  ax.set_xticks(list(range(len(time_tick_labels))))
  ax.set_xticklabels(time_tick_labels)
  if title is not None:
    ax.set_title(title)


def plot_timed_y_2(df, mult, signal_idx, fig, gs, title=None):
  windows = ['1s', '2s', '5s', '10s', '30s', '1m', '5m', '10m', '30m', '60m']
  timed_Y = []
  time_ticks = []
  for ws in reversed(windows):
    y = (df['bitmex-xbtusd/mid_ret.%s' % ws] / df['bitmex-xbtusd/b.mid_p'] * mult).iloc[signal_idx]
    timed_Y.append(y.mean())
    time_ticks.append('-%s' % ws)
  timed_Y.append(0.)
  time_ticks.append('0')
  for ws in windows:
    y = (df['Y/bitmex-xbtusd/mid_ret.%s.Y%s' %
            (ws, ws)] / df['bitmex-xbtusd/b.mid_p'] * mult).iloc[signal_idx]
    timed_Y.append(y.mean())
    time_ticks.append(ws)
  plot_timed_y(timed_Y, time_ticks, fig, gs, title=title)


def plot_res(buy_thr, sell_thr, Y_pred, Y_actual, df, fig, gs):
  gs = gs.subgridspec(3, 2)

  buy_mask = (buy_thr[0] < Y_pred) & (Y_pred < buy_thr[1])
  sell_mask = (sell_thr[0] < Y_pred) & (Y_pred < sell_thr[1])

  buy_idx = np.where(buy_mask)[0]
  sell_idx = np.where(sell_mask)[0]
  signal_idx = sorted(buy_idx.tolist() + sell_idx.tolist())

  print('thr B %f - %f, #=%d, Y-mean=%f' %
        (buy_thr[0], buy_thr[1], buy_idx.shape[0], Y_actual[buy_mask].mean()))
  print('thr S %f - %f, #=%d, Y-mean=%f' %
        (sell_thr[0], sell_thr[1], sell_idx.shape[0], Y_actual[sell_mask].mean()))

  mult = pd.Series(0., index=df.index)
  mult.iloc[buy_idx] = 1.
  mult.iloc[sell_idx] = -1.

  plot_bucket(Y_pred, Y_actual, fig, gs[:, 0])

  plot_timed_y_2(df, mult, signal_idx, fig, gs[0, 1], 'Timed-Y')
  plot_timed_y_2(df, mult, buy_idx, fig, gs[1, 1], 'Timed-Y Buy')
  plot_timed_y_2(df, mult, sell_idx, fig, gs[2, 1], 'Timed-Y Sell')


'''
def weired_fit(X, Y, ws=600):
  X = sm.add_constant(X)
  mod = sm.GLS(Y, X)
  res = X.fit()

  Y_pred = res.predict(X)
  sqr_error = (Y_pred - Y)**2

  for idx in range(0, X.shape[0], ws):
'''


def pipeline(X, df, os_ratio=0.3):
  # Dropping NA
  mask = (df['bitmex-xbtusd/b.mid_p'].rolling(2 * 60).std() > 4.6)
  X = X[mask]
  df = df[mask]

  cut_index = int(X.shape[0] * (1 - os_ratio))
  df_is, df_os = df.iloc[:cut_index], df.iloc[cut_index:]
  X_is, Y_is = X.iloc[:cut_index], df_is['Y']
  X_os, Y_os = X.iloc[cut_index:], df_os['Y']

  # Dropping NA
  mask_is = ((-20 <= X_is) & (X_is <= 20)).all(axis=1)
  X_is = X_is[mask_is]
  Y_is = Y_is[mask_is]
  df_is = df_is[mask_is]
  mask_os = ((-20 <= X_os) & (X_os <= 20)).all(axis=1)
  X_os = X_os[mask_os]
  Y_os = Y_os[mask_os]
  df_os = df_os[mask_os]

  # ----- Fit begin
  X2 = sm.add_constant(X_is)
  mod = sm.GLS(Y_is, X2)
  res = mod.fit()

  try:
    print(res.summary())
  except NotImplementedError:
    pass

  Y_is_pred = res.predict(X2)
  # (Y_pred - Y_is).to_pickle('dump.pkl')

  Y_pred_sorted = np.sort(Y_is_pred[~np.isnan(Y_is_pred)])
  size = Y_pred_sorted.shape[0]
  buy_thr = (Y_pred_sorted[int(0.85 * size)], Y_pred_sorted[int(0.95 * size)])
  sell_thr = (Y_pred_sorted[int(0.05 * size)], Y_pred_sorted[int(0.15 * size)])
  # ----- Fit end

  fig = plt.figure(figsize=(20, 10), constrained_layout=True)
  gs = gridspec.GridSpec(20, 2, figure=fig)

  # In-sample
  ax = fig.add_subplot(gs[0, 0])
  ax.axis('off')
  ax.text(0.1,
          0.5,
          'In-sample',
          horizontalalignment='center',
          verticalalignment='center',
          fontsize='large')
  print('In-sample')
  plot_res(buy_thr, sell_thr, Y_is_pred, Y_is, df_is, fig, gs[1:, 0])

  # Out-sample
  X2 = sm.add_constant(X_os)
  Y_os_pred = res.predict(X2)

  ax = fig.add_subplot(gs[0, 1])
  ax.axis('off')
  ax.text(0.1,
          0.5,
          'Out-sample',
          horizontalalignment='center',
          verticalalignment='center',
          fontsize='large')
  print('Out-sample')
  plot_res(buy_thr, sell_thr, Y_os_pred, Y_os, df_os, fig, gs[1:, 1])

  plt.savefig('test.png')


def main(argv):
  setup_plt()
  df = pd.read_hdf('dump_concat.hd5.lz4')

  ref_symbols = [
      'bitflyer-fxbtcjpy',  # 'bitmex-xbtusd',
      'huobi-btcusd-quarter',
      'huobi-btcusd-next_week',
      'huobi-btcusd-this_week',
      'kraken-btcusd-perpetual',
      'okex-btcusd-quarter',
      'okex-btcusd-next_week',
      'okex-btcusd-this_week',
      'okex-btcusd-perpetual',
      'binance-btcusdt',
      'bitstamp-btcusd',
      'gdax-btcusd',
      'huobi-btcusdt',
      'kraken-btcusd',
      'okex_btcusdt',
      'bitflyer-btcjpy',
      'bithumb-btckrw',
      'upbit-btckrw'
  ]

  if not os.path.exists('X.hd5.lz4'):
    X = pd.DataFrame(index=df.index)
    for ref_symbol in ref_symbols:
      for w in [1, 2, 5, 10, 15, 20, 30, 60, 120]:
        ref_p = (
            df['%s/b.mid_p' % ref_symbol] +
            (df['bitmex-xbtusd/b.mid_p'] - df['%s/b.mid_p' % ref_symbol]).rolling(w * 60).mean())
        X['%s/w%d' %
          (ref_symbol, w)] = (ref_p - df['bitmex-xbtusd/b.mid_p']) / df['bitmex-xbtusd/b.mid_p']
    X.to_hdf('X.hd5.lz4', 'dataframe', complib='blosc:lz4')
  else:
    X = pd.read_hdf('X.hd5.lz4')
  print(X.shape)

  df['Y/bitmex-xbtusd/mid_ret.60m.Y60m'] = df['Y/bitmex-xbtusd/mid_ret.60m.Y1h']
  df['Y'] = df['Y/bitmex-xbtusd/mid_ret.5m.Y5m'] / df['bitmex-xbtusd/b.mid_p']

  pipeline(X, df)
  # plt.figure(num=None, figsize=(7.5, 4.5))
  # plt.savefig('test.png')


if __name__ == '__main__':
  logging.basicConfig(
      level='DEBUG',
      format='%(levelname)8s %(asctime)s %(name)s %(filename)s:%(lineno)d] %(message)s')
  app.run(main)
